# -*- coding: utf-8 -*-

import math
import heapq
import numpy as np
import matplotlib.pyplot as plt

# ================================
# 参数区（根据需要调整）
# ================================
SEED = 42

# 飞行器运动学 & 规划参数
MAX_CLIMB_DEG = 25.0      # 最大爬升/俯冲角（绝对值）建议10~30°
STEP_LEN = 2.0            # 单步前进长度 s，越小越稳定但更慢
YAW_DISCRETIZATION_DEG = 20.0   # 每步偏航增量离散步（±这个角度，含0）
PITCH_BLEND = 0.5         # 俯仰角平滑系数 0.3~0.6
A_STAR_WEIGHT = 1.5       # A*权重 self.w
VERTICAL_WEIGHT_ALPHA = 2.0  # 启发式对竖直距离加权 α（1.5~3.0）
LAMBDA_Z = 0.05           # “高度进度偏差”惩罚系数 0.02~0.1

# 分段爬升→水平巡航形状参数
CLIMB_END_FRAC = 0.30     # 在水平进度的多少时达到目标高度（0.5=一半，0.6=60%）
ALPHA_LN = 1.0            # ln(1+αx) 曲线形状，越大前期越陡
CRUISE_DEADBAND = 0.2     # 巡航阶段对高度误差的“死区”米，避免抖动

# 障碍与安全距离（如无障碍，ignore_obstacles=True）
IGNORE_OBSTACLES = True
SAFE_DIS = 0.5

# 地图尺寸（如你已有地图，请替换 WorldMap 构建逻辑）
WORLD_MIN = np.array([0.0, 0.0, 0.0])
WORLD_MAX = np.array([20000.0, 20000.0, 400.0])
WORLD_RES = 1.0  # 仅用于离散检查时的索引换算

# 多机起终点示例
UAV_TASKS = [
    # (start(x,y,z,yaw_deg), goal(x,y,z))
    ((10.0, 10.0, 2.0, 0.0),   (5000.0, 5000.0, 150.0)),
    ((100.0, 100.0, 2.0, 0.0),   (15000.0, 16000.0, 180.0)),
    ((210.0, 10.0, 2.0, 0.0),   (1000.0, 10000.0, 250.0)),
    ((300.0, 300.0, 2.0, 0.0),   (2000.0, 16000.0, 180.0)),
    ((400.0, 400.0, 2.0, 0.0),   (15000.0, 5000.0, 200.0)),
    ((500.0, 100.0, 2.0, 0.0),   (16000.0, 15000.0, 280.0)),
    ((100.0, 500.0, 2.0, 0.0),   (12000.0, 9000.0, 280.0)),
    ((1000.0, 500.0, 2.0, 0.0),   (17000.0, 8000.0, 280.0)),
    ((1500.0, 500.0, 2.0, 0.0),   (12000.0, 10000.0, 280.0)),
    ((2000.0, 500.0, 2.0, 0.0),   (15000.0, 6000.0, 280.0)),
]

# 随机障碍示例开关（仅演示；真实项目中请用你的 ESDF/占据图）
ADD_DEMO_OBS = False
NUM_OBS = 0                # 随机障碍数量（演示用）

# ================================
# 工具函数
# ================================
def deg2rad(d): return d * math.pi / 180.0
def rad2deg(r): return r * 180.0 / math.pi

def clamp(v, lo, hi): return max(lo, min(hi, v))

# ================================
# 数据结构
# ================================
class ContinuousNode:
    __slots__ = (
        "x", "y", "z", "pitch", "yaw", "g", "f", "parent", "velocity", "key"
    )
    def __init__(self, x, y, z, pitch, yaw, g, f, parent, velocity):
        self.x = x
        self.y = y
        self.z = z
        self.pitch = pitch
        self.yaw = yaw
        self.g = g
        self.f = f
        self.parent = parent
        self.velocity = velocity
        # 键用于哈希/去重（可按需粗化）
        self.key = (round(x,2), round(y,2), round(z,2), round(rad2deg(yaw),1), round(rad2deg(pitch),1))

class Kinematics:
    def __init__(self, max_climb_deg=15.0, velocity=10.0):
        self.max_climb_angle = deg2rad(max_climb_deg)
        self.velocity = velocity  # m/s（仅计时用）       
# ================================
# 世界/地图（占位实现）
# ================================
class WorldMap:
    def __init__(self, min_xyz, max_xyz, res=1.0, add_demo_obs=False):
        self.min_x, self.min_y, self.min_z = min_xyz
        self.max_x, self.max_y, self.max_z = max_xyz
        self.resolution = res

        self.x_width = int(round((self.max_x - self.min_x) / self.resolution)) + 1
        self.y_width = int(round((self.max_y - self.min_y) / self.resolution)) + 1
        self.z_width = int(round((self.max_z - self.min_z) / self.resolution)) + 1

        # 如果你有真实的 ESDF/占据图，请在此加载：
        self.esdf_map = None
        self.obs_map = None

        if add_demo_obs:
            # 演示随机障碍（稀疏点）
            self.obs_map = np.zeros((self.x_width, self.y_width, self.z_width), dtype=bool)
            rng = np.random.default_rng(SEED)
            for _ in range(NUM_OBS):
                ix = rng.integers(0, self.x_width)
                iy = rng.integers(0, self.y_width)
                iz = rng.integers(0, self.z_width)
                self.obs_map[ix, iy, iz] = True

# ================================
# A* 连续固定翼
# ================================
class ContinuousAstarFixedWing:
    def __init__(self, world_map: WorldMap, kin: Kinematics,
                 step_len=2.0, yaw_discretization_deg=20.0,
                 w=1.5, safe_dis=0.5, ignore_obstacles=True,
                 lambda_z=0.05, alpha_vert=2.0, pitch_blend=0.5,
                 climb_end_frac=CLIMB_END_FRAC, alpha_ln=ALPHA_LN, cruise_deadband=CRUISE_DEADBAND):
        self.map = {
            "min_x": world_map.min_x, "min_y": world_map.min_y, "min_z": world_map.min_z,
            "max_x": world_map.max_x, "max_y": world_map.max_y, "max_z": world_map.max_z,
            "resolution": world_map.resolution,
            "x_width": world_map.x_width, "y_width": world_map.y_width, "z_width": world_map.z_width,
            "esdf_map": world_map.esdf_map, "obs_map": world_map.obs_map
        }
        self.kin = kin
        self.s = step_len
        self.w = w
        self.safe_dis = safe_dis
        self.ignore_obstacles = ignore_obstacles
        self.lambda_z = lambda_z
        self.alpha_vert = alpha_vert
        self.pitch_blend = pitch_blend

        # 分段爬升→水平巡航参数
        self.climb_end_frac = clamp(climb_end_frac, 1e-3, 0.99)  # 0<frac<1
        self.alpha_ln = max(1e-6, alpha_ln)
        self.cruise_deadband = max(0.0, cruise_deadband)

        dy = deg2rad(yaw_discretization_deg)
        # 原始实现常用的 primitives 格式 (s, dyaw, dp_unused)
        self.primitives = [(self.s, 0.0, 0.0), (self.s, +dy, 0.0), (self.s, -dy, 0.0)]

        # 运行时变量
        self.start = None
        self.xy_total = 1.0  # 水平总距离（用于进度）
        self.counter = 0
    # ---------- 基本工具 ----------
    def _inside_world(self, x, y, z):
        return (self.map["min_x"] <= x <= self.map["max_x"] and
                self.map["min_y"] <= y <= self.map["max_y"] and
                self.map["min_z"] <= z <= self.map["max_z"])

    def heuristic(self, node: ContinuousNode, goal: ContinuousNode):
        dx = goal.x - node.x
        dy = goal.y - node.y
        dz = goal.z - node.z
        # 对竖直距离加权（温和地鼓励提前满足高度）
        return math.sqrt(dx*dx + dy*dy + self.alpha_vert*dz*dz)

    # ---------- 扩展：强约束版（分段 ln 爬升→水平巡航） ----------
    def _expand(self, node: ContinuousNode, goal: ContinuousNode):
        succ = []

        # 当前到 goal 的水平距离
        horiz_to_goal = math.hypot(goal.x - node.x, goal.y - node.y)
        # 已完成的水平进度（0~1）
        progress = 1.0 - horiz_to_goal / self.xy_total
        progress = clamp(progress, 0.0, 1.0)

        # 分段参考高度：前段按 ln(1+αx) 快速爬升，后段等高水平
        p_end = self.climb_end_frac           # 例如 0.5 或 0.6
        dz_goal = (goal.z - self.start.z)

        if progress < p_end:
            # 归一到 [0,1] 的“爬升子进度”
            t = progress / max(1e-6, p_end)
            # ln 型（前快后缓），严格归一化到 [0,1]
            f = math.log(1.0 + self.alpha_ln * t) / math.log(1.0 + self.alpha_ln)
            z_ref = self.start.z + f * dz_goal
        else:
            # 到达爬升终点后保持巡航高度（等高飞行）
            z_ref = goal.z

        for s, dyaw, _ in self.primitives:
            new_yaw = node.yaw + dyaw

            # 以“使本步后靠近 z_ref”为目标反推期望 pitch
            dz_need = z_ref - node.z

            # 巡航阶段：在小死区内强制 0，避免抖动
            if progress >= p_end and abs(dz_need) < self.cruise_deadband:
                dz_need = 0.0

            desired_pitch = math.atan2(dz_need, max(1e-6, s))

            # 平滑并限幅到飞行器能力
            new_pitch = node.pitch + (desired_pitch - node.pitch) * self.pitch_blend
            new_pitch = clamp(new_pitch, -self.kin.max_climb_angle, self.kin.max_climb_angle)

            # 前进一步
            horiz = s * math.cos(new_pitch)
            dx = horiz * math.cos(new_yaw)
            dy = horiz * math.sin(new_yaw)
            dz = s * math.sin(new_pitch)

            nx, ny, nz = node.x + dx, node.y + dy, node.z + dz

            if not self._inside_world(nx, ny, nz):
                continue

            # 碰撞/安全检查
            if not self.ignore_obstacles:
                ix = int(round((nx - self.map['min_x']) / self.map['resolution']))
                iy = int(round((ny - self.map['min_y']) / self.map['resolution']))
                iz = int(round((nz - self.map['min_z']) / self.map['resolution']))
                if ix < 0 or iy < 0 or iz < 0 or ix >= self.map['x_width'] or iy >= self.map['y_width'] or iz >= self.map['z_width']:
                    continue
                esdf = self.map.get('esdf_map', None)
                if esdf is None:
                    occ = self.map.get('obs_map', None)
                    if occ is not None and occ[ix, iy, iz]:
                        continue
                else:
                    if esdf[ix, iy, iz] < self.safe_dis:
                        continue

            # g 代价：时间 + 高度进度偏差惩罚
            g_time = node.g + s / self.kin.velocity
            z_err = abs(nz - z_ref)
            g_new = g_time + self.lambda_z * z_err

            child = ContinuousNode(nx, ny, nz, new_pitch, new_yaw, g_new, 0.0, node, self.kin.velocity)
            child.f = child.g + self.w * self.heuristic(child, goal)
            succ.append(child)

        return succ

    # ---------- 规划 ----------
    def plan(self, start_xyz_yaw, goal_xyz, max_iter=200000):
        sx, sy, sz, syaw_deg = start_xyz_yaw
        gx, gy, gz = goal_xyz

        self.start = ContinuousNode(sx, sy, sz, 0.0, deg2rad(syaw_deg), 0.0, 0.0, None, self.kin.velocity)
        goal = ContinuousNode(gx, gy, gz, 0.0, 0.0, 0.0, 0.0, None, self.kin.velocity)

        # 计算“水平总距离”用于进度
        self.xy_total = max(1e-6, math.hypot(gx - sx, gy - sy))

        open_heap = []
        heapq.heappush(open_heap, (0.0, self.counter, self.start))
        self.counter += 1

        open_dict = {self.start.key: self.start}
        closed = set()

        goal_reached = None
        iter_count = 0

        # 收敛阈值（到达半径）——可按需调整
        reach_tol_xy = max(self.s * 1.5, 2.0)
        reach_tol_z = max(self.s * 0.8, 1.0)

        while open_heap and iter_count < max_iter:
            iter_count += 1
            _,_, cur = heapq.heappop(open_heap)

            if cur.key in closed:
                continue
            closed.add(cur.key)

            # 判断到达：水平与高度均接近
            if (abs(cur.x - gx) <= reach_tol_xy and
                abs(cur.y - gy) <= reach_tol_xy and
                abs(cur.z - gz) <= reach_tol_z):
                goal_reached = cur
                break

            # 扩展
            for nxt in self._expand(cur, goal):
                if nxt.key in closed:
                    continue
                old = open_dict.get(nxt.key, None)
                if old is None or nxt.g < old.g:
                    open_dict[nxt.key] = nxt
                    heapq.heappush(open_heap, (nxt.f, self.counter, nxt))
                    self.counter += 1


        if goal_reached is None:
            return None  # 未找到

        # 回溯路径
        path = []
        n = goal_reached
        while n is not None:
            path.append((n.x, n.y, n.z))
            n = n.parent
        path.reverse()
        return path


# ================================
# 演示主程序（多机顺序规划）
# ================================
def run_multi_planning():
    print("Running multi continuous planning for {} aircraft...".format(len(UAV_TASKS)))
    world = WorldMap(WORLD_MIN, WORLD_MAX, WORLD_RES, add_demo_obs=ADD_DEMO_OBS)
    kin = Kinematics(max_climb_deg=MAX_CLIMB_DEG, velocity=12.0)  # 速度仅用于时间成本

    planner = ContinuousAstarFixedWing(
        world_map=world, kin=kin,
        step_len=STEP_LEN,
        yaw_discretization_deg=YAW_DISCRETIZATION_DEG,
        w=A_STAR_WEIGHT,
        safe_dis=SAFE_DIS,
        ignore_obstacles=IGNORE_OBSTACLES,
        lambda_z=LAMBDA_Z,
        alpha_vert=VERTICAL_WEIGHT_ALPHA,
        pitch_blend=PITCH_BLEND,
        climb_end_frac=CLIMB_END_FRAC,   # 在进度 60% 达到目标高度
        alpha_ln=ALPHA_LN,               # ln 曲线形状
        cruise_deadband=CRUISE_DEADBAND  # 巡航死区，抑制抖动
    )

    all_paths = []
    fig = plt.figure(figsize=(9, 6))
    ax = fig.add_subplot(111, projection='3d')
    ax.set_xlim(WORLD_MIN[0], WORLD_MAX[0])
    ax.set_ylim(WORLD_MIN[1], WORLD_MAX[1])
    ax.set_zlim(WORLD_MIN[2], WORLD_MAX[2])
    ax.set_xlabel('X [m]')
    ax.set_ylabel('Y [m]')
    ax.set_zlabel('Z [m]')

    colors = ['tab:blue', 'tab:orange', 'tab:green', 'tab:red']
    any_success = False

    for i, (start, goal) in enumerate(UAV_TASKS, start=1):
        path = planner.plan(start, goal)
        if path is None or len(path) < 2:
            print(f"Task failed: No path found")
            print(f"UAV{i}: planning failed")
            # 画起终点，留空轨迹
            sx, sy, sz, _ = start
            gx, gy, gz = goal
            ax.scatter([sx], [sy], [sz], marker='o', s=40, label=f'UAV{i} start')
            ax.scatter([gx], [gy], [gz], marker='^', s=50, label=f'UAV{i} goal')
            continue

        any_success = True
        all_paths.append(path)
        pts = np.array(path)
        sx, sy, sz, _ = start
        gx, gy, gz = goal

        c = colors[(i-1) % len(colors)]
        ax.plot(pts[:,0], pts[:,1], pts[:,2], lw=2.0, label=f'UAV{i} path', color=c)
        ax.scatter([sx], [sy], [sz], marker='o', s=40, label=f'UAV{i} start', color=c)
        ax.scatter([gx], [gy], [gz], marker='^', s=50, label=f'UAV{i} goal', color=c)

        print(f"UAV{i}: path length = {len(path)} steps")

    # 将图例放到画布外侧，避免遮挡轨迹；使用 tight bbox 保存时包含图例
    handles, labels = ax.get_legend_handles_labels()
    if len(handles) > 0:
        # 把图例放在右侧画布外，留出右侧空间
        ax.legend(loc='upper left', bbox_to_anchor=(1.02, 1), borderaxespad=0.)

    # 在子图布局中为右侧图例保留空间，然后保存时使用 bbox_inches='tight' 避免裁剪
    plt.tight_layout(rect=[0, 0, 0.85, 1.0])
    plt.savefig("multi_continuous_paths.png", dpi=180, bbox_inches='tight')
    print("Saved plot to multi_continuous_paths.png")
    plt.show()

    return any_success


if __name__ == "__main__":
    ok = run_multi_planning()
